import os

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import os


class Config(object):
    def __init__(self, dataset, embedding):
        self.model_name = 'TextCNN_capital_loan'
        self.train_path = dataset + '/data/train.txt'
        self.dev_path = dataset + '/data/dev.txt'
        self.test_path = dataset + '/data/test.txt'
        print(os.getcwd())
        self.class_list = [x.strip() for x in open(
            dataset + '/data/class.txt', encoding='utf-8'
        ).readlines()]
        self.vocab_path = dataset + '/data/vocab_capital.pkl'
        self.save_path = dataset + '/saved_dict/' + self.model_name + '.ckpt'
        self.log_path = dataset + '/log/' + self.model_name
        self.embedding_pretrained = torch.tensor(
            np.load(dataset + '/data/' + embedding)['embeddings'].astype('float32')
        ) if embedding != 'random' else None

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.dropout = 0.5
        self.require_improvement = 1000
        self.num_classes = len(self.class_list)
        self.n_vocab = 0
        self.num_epochs = 20
        self.batch_size = 128
        self.pad_size = 43
        self.learning_rate = 1e-3
        self.embed = self.embedding_pretrained.size(1) \
            if self.embedding_pretrained is not None else 300
        self.filter_sizes = (2, 3, 4)
        self.num_filters = 256


class Model(nn.Module):
    def __init__(self,config):
        super(Model, self).__init__()
        if config.embedding_pretrained is not None:
            self.embedding = nn.Embedding.from_pretrained(config.embedding_pretrained,freeze=False)
        else:
            self.embedding = nn.Embedding(config.n_vocab, config.embed, padding_idx=config.n_vocab - 1) #词嵌入层
        self.convs = nn.ModuleList(
            [nn.Conv2d(1,config.num_filters,(k,config.embed)) for k in config.filter_sizes] #网络层数
        )
        self.dropout = nn.Dropout(config.dropout) #防止过拟合
        self.fc = nn.Linear(config.num_filters * len(config.filter_sizes),config.num_classes)



    def conv_and_pool(self,x,conv):
        x = F.relu(conv(x).squeeze(3))
        x = F.max_pool1d(x,x.size(2)).squeeze(2)
        return x

    def forward(self,x):
        out = self.embedding(x[0])   #要知道X是长什么样子的
        out = out.unsqueeze(1)
        out = torch.cat([self.conv_and_pool(out,conv) for conv in self.convs],1)
        out = self.dropout(out)
        out = self.fc(out)
        return out
